# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 3
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####

## VP module
## Copyright (c) 2012 by Christopher Koch



"""This module contains classes and methods for handling "Volition Package" (VP) files."""

## No guarantees about pep8 compliance

import time
from struct import *


class VolitionError(Exception):
    """Base class for exceptions in this module."""
    pass
    

class FileFormatError(VolitionError):
    """Exception raised for invalid filetype errors.

    Attributes:
        path -- the filepath of the invalid file
        msg  -- a message"""
    def __init__(self, path, msg):
        self.path = path
        self.msg = msg
        
    def __str__(self):
        return "{}, {}.".format(self.msg, self.path)
        

class FileNotFoundError(VolitionError):
    def __init__(self, path, msg):
        self.path = path
        self.msg = msg

    def __str__(self):
        return "File not found: {}, {}.".format(self.path, self.msg)
        

class VolitionPackageError(VolitionError):
    def __init__(self, msg):
        self.msg = msg

    def __str__(self):
        return self.msg
        

# Note for the pack() and unpack() wrappers:
# u = Unpacked data
# p = Packed data
# Mind your p's and u's

def unpack_byte(bin_data):
    """Wrapper function for struct.unpack().  Can accept an iterable of any length and will unpack the contents into a list of integers."""

    u = int()
    try:
        p = bytes(bin_data)
    except TypeError:
        p = bytes(bin_data, "utf-8", "ignore")

    if len(p) == 1:
        u = unpack('b', p)[0]

    elif len(p) > 1:
        u = list(unpack('{}b'.format(len(p)), p))

    return u
    

def unpack_ubyte(bin_data):

    # unsigned byte (numeric)

    u = int()
    try:
        p = bytes(bin_data)
    except TypeError:
        p = bytes(bin_data, "utf-8", "ignore")

    if len(p) == 1:
        u = unpack('B', p)[0]

    elif len(p) > 1:
        u = list(unpack('{}B'.format(len(p)), p))

    return u
    

def unpack_short(bin_data):

    # signed short

    u = int()
    try:
        p = bytes(bin_data)
    except TypeError:
        p = bytes(bin_data, "utf-8", "ignore")

    if len(p) == 2:
        u = unpack('h', p)[0]

    elif len(p) > 2 and (len(p) % 2) == 0:
        u = list(unpack('{}h'.format(len(p) / 2), p))

    return u
    

def unpack_ushort(bin_data):

    # unsigned short

    u = int()
    try:
        p = bytes(bin_data)
    except TypeError:
        p = bytes(bin_data, "utf-8", "ignore")

    if len(p) == 2:
        u = unpack('H',p)[0]

    elif len(p) > 2 and (len(p) % 2) == 0:
        u = list(unpack('{}H'.format(len(p) / 2), p))

    return u
    

def unpack_int(bin_data):

    # signed int32

    u = int()
    try:
        p = bytes(bin_data)
    except TypeError:
        p = bytes(bin_data, "utf-8", "ignore")

    if len(p) == 4:
        u = unpack('i', p)[0]

    elif len(p) > 4 and (len(p) % 4) == 0:
        u = list(unpack('{}i'.format(len(p) // 4), p))

    return u
    

def unpack_uint(bin_data):

    # unsigned int32

    u = int()
    try:
        p = bytes(bin_data)
    except TypeError:
        p = bytes(bin_data, "utf-8", "ignore")

    if len(p) == 4:
        u = unpack('I', p)[0]

    elif len(p) > 4 and (len(p) % 4) == 0:
        u = list(unpack('{}I'.format(len(p) / 4), p))

    return u
    

def unpack_float(bin_data):

    # float

    u = float()
    try:
        p = bytes(bin_data)
    except TypeError:
        p = bytes(bin_data, "utf-8", "ignore")

    if len(p) == 4:
        u = unpack('f', p)[0]

    elif len(p) > 4 and (len(p) % 4) == 0:
        u = list(unpack('{}f'.format(len(p) / 4), p))

    return u
    

def unpack_vector(bin_data):

    # tuple of three floats

    u = tuple()
    try:
        p = bytes(bin_data)
    except TypeError:
        p = bytes(bin_data, "utf-8", "ignore")

    #print(len(p))

    if len(p) == 12:
        u = unpack('3f', p)

    elif len(p) > 12 and (len(p) % 12) == 0:
        u = list()
        for i in range(len(p) // 12):
            u.append(unpack('3f', p[i * 12:i * 12 + 12]))

    return tuple(u)
    

def pack_byte(x):

    # signed byte

    p = bytes()
    try:
        u = tuple(x)
        p = pack('b', u[0])
        for i in u[1:]:
            p += pack('b', i)
    except TypeError:
        p = pack('b', x)

    return p
    

def pack_ubyte(x):

    # unsigned byte

    p = bytes()
    try:
        u = tuple(x)
        p = pack('B', u[0])
        for i in u[1:]:
            p += pack('b', i)
    except TypeError:
        p = pack('B', x)

    return p
    

def pack_short(x):

    # signed short

    p = bytes()
    try:
        u = tuple(x)
        p = pack('h', u[0])
        for i in u[1:]:
            p += pack('h', i)
    except TypeError:
        p = pack('h', x)

    return p
    

def pack_ushort(x):

    # unsigned short

    p = bytes()
    try:
        u = tuple(x)
        p = pack('H', u[0])
        for i in u[1:]:
            p += pack('H', i)
    except TypeError:
        p = pack('H', x)

    return p
    

def pack_int(x):

    # signed int32

    p = bytes()
    try:
        u = tuple(x)
        p = pack('i', u[0])
        for i in u[1:]:
            p += pack('i', i)
    except TypeError:
        p = pack('i', x)

    return p
    

def pack_uint(x):

    # unsigned int32

    p = bytes()
    try:
        u = tuple(x)
        p = pack('I', u[0])
        for i in u[1:]:
            p += pack('I', i)
    except TypeError:
        p = pack('I', x)

    return p
    

def pack_float(x):

    # float

    p = bytes()
    try:
        u = tuple(x)
        p = pack('f', u[0])
        for i in u[1:]:
            p += pack('f', i)
    except TypeError:
        p = pack('f', x)

    return p
    

def pack_string(x):

    # int with length of string followed by chars

    u = bytes(x)
    p = pack('i', len(u))
    p += u

    return p
    

class RawData:
    """Creates an object that can be read like a file.  Takes any sequence of data as an argument.  May typically be used to pass only a part of a file to a function so that the part of the file can still be read like a file.

    Methods:
        read(length = 0) -- Returns a slice of data from the current address to the current address plus length.  Increases current address by length.  If length is 0 or not provided, returns the entire data.
        seek(new_addr[, whence = 0]) -- Changes the current address.  If whence is 0, new_addr is bytes from beginning; if whence is 1, new_addr is bytes from current address; if whence is 2, new_addr is bytes from end."""
    def __init__(self, data):
        self.data = data
        self.addr = 0

    def __len__(self):
        return len(self.data)

    def __repr__(self):
        return "<RawData object of length {} at index {}>".format(len(self.data), self.addr)

    def read(self, length=None):
        if length is None:
            return self.data
        elif length == 0:
            return b""
        else:
            out = self.data[self.addr:self.addr + length]
            self.addr += length
            return out

    def seek(self, new_addr, whence = 0):
        if whence == 1:
            self.addr += new_addr
        elif whence == 2:
            self.addr = len(self.data) - new_addr
        else:
            self.addr = new_addr

    def tell(self):
        return self.addr

class VolitionPackageFile:
    """A VP directory tree.  Folders are stored as Folder objects and files are stored as File objects containing
    the contents of the file.  The constructor takes a file-like object as a required argument and creates the
    necessary File and Folder objects to fill the directory tree, self.vp_file_directory.  Internally, all folders
    and files are stored in a Folder object with the name attribute "root."  However, all read and write methods
    are written to take paths starting with the first entry in the actual VP file's index (usually "data")."""
    def __init__(self, vp_file):
        vp_file_id = vp_file.read(4)
        if vp_file_id != b"VPVP":
            raise FileFormatError(vp_file, "Not a VP file.")

        vp_file_version = unpack_int(vp_file.read(4))
        if vp_file_version > 2:
            raise FileFormatError(vp_file, "VP spec version greater than 2.")

        vp_diroffset = unpack_int(vp_file.read(4))
        vp_num_files = unpack_int(vp_file.read(4))

        vp_file_directory = Folder(b"root")
        parent_directory = vp_file_directory

        vp_file.seek(vp_diroffset)

        for i in range(vp_num_files):
            this_file_offset = unpack_int(vp_file.read(4))
            this_file_size = unpack_int(vp_file.read(4))
            this_file_name_long = vp_file.read(32)
            this_fname_len = this_file_name_long.index(b"\0")
            #this_file_name = this_file_name_long.rstrip(b'\0')
            this_file_name = this_file_name_long[:this_fname_len]
            this_file_timestamp = unpack_int(vp_file.read(4))
            vp_dir_location = vp_file.tell()

            if not this_file_offset or not this_file_size or not this_file_timestamp:
                # direntry is either folder or backdir

                if this_file_name != b"..":     # folder
                    this_node = Folder(this_file_name, parent_directory)
                    parent_directory.contents.add(this_node)
                    parent_directory = this_node
                else:                           # backdir
                    this_node = parent_directory
                    parent_directory = this_node.parent
            else:
                # direntry is a file

                vp_file.seek(this_file_offset)
                this_file_data = vp_file.read(this_file_size)
                this_node = File(this_file_name, this_file_data, this_file_timestamp, parent_directory)
                parent_directory.contents.add(this_node)
                vp_file.seek(vp_dir_location)

        self.vp_file_directory = vp_file_directory

    def get_file(self, path, sep='/'):
        """Returns a File object or Folder object for further processing.  Takes a full path including the file/folder name
        as a required argument and a separator character as an optional argument."""
        parent_directory = self.vp_file_directory
        split_path = path.split(sep)
        cur_node = parent_directory.contents

        # traverse path
        for i, cur_dir in enumerate(split_path):
            for dir in cur_node:
                if dir.name == cur_dir:
                    parent_directory = dir
                    cur_node = dir.contents
                    break
            else:
                raise FileNotFoundError(path, "{} does not exist".format(cur_dir))

        # At this point, dir or parent_directory matches the
        # last item in split_path and cur_node is the contents
        # of that item, whether it's a file or folder.
        # parent_directory is a File object or Folder object
        # cur_node is a set of File objects and/or Folder objects

        return parent_directory

    def remove_file(self, path, sep='/'):
        """Removes a file or folder from the directory.  Takes a full path including the file/folder name
        as a required argument and a separator character as an optional argument."""
        parent_directory = self.vp_file_directory
        split_path = path.split(sep)
        cur_node = parent_directory.contents
        path_depth = len(split_path) - 1

        # traverse path
        for i, cur_dir in enumerate(split_path):
            for dir in cur_node:
                if dir.name == split_path[i]:
                    parents_parent = parent_directory
                    parent_directory = dir
                    cur_node = dir.contents
                    break
            else:
                raise FileNotFoundError(path, "{} does not exist".format(cur_dir))

        # At this point, dir or parent_directory matches the
        # last item in split_path and cur_node is the contents
        # of that item, whether it's a file or folder.
        # parent_directory is a File object or Folder object
        # cur_node is a set of File objects and/or Folder objects

        if parent_directory.name == split_path[path_depth]:
            parents_parent.contents.remove(parent_directory)
        else:
            raise FileNotFoundError(path, " does not exist")

    def add_file(self, path, file, sep='/'):
        parent_directory = self.vp_file_directory
        split_path = path.split(sep)

        cur_node = parent_directory.contents

        path_depth = len(split_path)

        for i, cur_dir in enumerate(split_path):
            for dir in cur_node:
                if dir.name == cur_dir:
                    parent_directory = dir
                    cur_node = dir.contents
                    break
            else:
                raise FileNotFoundError(path, "{} does not exist".format(cur_dir))

        # At this point, dir or parent_directory matches the
        # last item in split_path and cur_node is the contents
        # of that item, whether it's a file or folder.
        # parent_directory is a File object or Folder object
        # cur_node is a set of File objects and/or Folder objects

        cur_node.add(file)

    def make_vp_file(self):
        self.cur_index = list()
        self.cur_files = list()
        self.cur_offset = 16        # header size == beginning offset is always 16 bytes
        self._recurse_thru_directory(self.vp_file_directory.contents)

        vp_dir_offset = self.cur_offset
        vp_num_files = len(self.cur_index)

        vp_header = b"".join([b"VPVP",
                             pack_int(2),
                             pack_int(vp_dir_offset),
                             pack_int(vp_num_files)])
        vp_files = b"".join(self.cur_files)
        vp_index = b"".join(self.cur_index)

        return b"".join([vp_header, vp_files, vp_index])

    def _recurse_thru_directory(self, vp_file_directory):
        # vp_file_directory should be a set - the contents of a Folder
        cur_index = self.cur_index
        cur_files = self.cur_files
        cur_offset = self.cur_offset

        # This function should recurse through vp_file_directory, not returning anything,
        # but updating self.cur_index, etc. at the end of each branch

        for cur_node in vp_file_directory:
            if isinstance(cur_node, Folder):
                this_entry_offset = 0
                this_entry_size = 0
                this_entry_name = cur_node.name.encode()
                this_entry_timestamp = 0

                # add index entry for folder:
                cur_index.append(b"".join([pack_int(this_entry_offset),
                                          pack_int(this_entry_size),
                                          this_entry_name.ljust(32, b'\0'),
                                          pack_int(this_entry_timestamp)]))

                # save current state:
                self.cur_index = cur_index
                self.cur_files = cur_files
                self.cur_offset = cur_offset

                self._recurse_thru_directory(cur_node.contents)

                # load current state:
                cur_index = self.cur_index
                cur_files = self.cur_files
                cur_offset = self.cur_offset
            elif isinstance(cur_node, File):
                this_entry_offset = cur_offset
                this_entry_size = len(cur_node)
                this_entry_name = cur_node.name.encode()
                this_entry_timestamp = cur_node.timestamp

                cur_offset += this_entry_size
                cur_files.append(cur_node.contents)

                # add index entry for file:
                cur_index.append(b"".join([pack_int(this_entry_offset),
                                          pack_int(this_entry_size),
                                          this_entry_name.ljust(32, b'\0'),
                                          pack_int(this_entry_timestamp)]))
            else:
                raise VolitionPackageError("{} is {}, expected File or Folder".format(cur_node, type(cur_node)))

        # add backdir
        this_entry_offset = 0
        this_entry_size = 0
        this_entry_name = b".."
        this_entry_timestamp = 0

        # add index entry:
        cur_index.append(b"".join([pack_int(this_entry_offset),
                                  pack_int(this_entry_size),
                                  this_entry_name.ljust(32, b'\0'),
                                  pack_int(this_entry_timestamp)]))

        # save current state:
        self.cur_index = cur_index
        self.cur_files = cur_files
        self.cur_offset = cur_offset

class Folder:
    def __init__(self, name, parent="", contents=None):
        self.name = name.decode(errors="ignore").lower()
        if contents is not None:
            self.contents = set(contents)
        else:
            self.contents = set()
        self.parent = parent
        self.parent_name = str(parent)

    def __eq__(self, other):
        if self.name == other.name and self.parent == other.parent:
            return True
        else:
            return False

    def __hash__(self):
        return hash(self.name)

    def __repr__(self):
        return "/".join([self.parent_name, self.name])

class File:
    def __init__(self, name, contents, timestamp=False, parent=""):
        try:
            self.name = name.decode().lower()
        except AttributeError:
            self.name = name
        self.contents = contents
        if timestamp:
            self.timestamp = timestamp
        else:
            self.timestamp = int(time.clock())
        self.parent = parent
        self.parent_name = str(parent)

    def __len__(self):
        return len(self.contents)

    def __eq__(self, other):
        if self.name == other.name and self.parent == other.parent:
            return True
        else:
            return False

    def __hash__(self):
        return hash(self.name)

    def __repr__(self):
        return "/".join([self.parent_name, self.name])